'''
作者：小鹿
个人网站：https://www.siena.zone
创作日期：2020年5月
'''
import cv2 as cv
import os

class WaterMark():

    def watermark(self, imgpath, markpath, pos_index, opacity, scale):
        self.image = cv.imread(imgpath)
        self.watermark_origin = cv.imread(markpath)
        if self.watermark_origin.shape[2] < 3:
            self.watermark = cv.cvtColor(self.watermark_origin, cv.COLOR_GRAY2BGR)
        else:
            self.watermark = self.watermark_origin
        self.bleedingWidth = scale / 20
        self.w, self.h = self.image.shape[1], self.image.shape[0]
        '''if self.w > self.h:
            self.rate = int(self.w * 0.1) / self.watermark.shape[1]
        else:
            self.rate = int(self.h * 0.1) / self.watermark.shape[0]'''
        self.rate = min(self.w, self.h) * scale / max(self.watermark.shape[1], self.watermark.shape[0])
        self.watermark_resize = cv.resize(self.watermark, None, fx=self.rate, fy=self.rate)

        #print(self.watermark_resize.shape[1])
        #print(self.watermark_resize.shape[0])
        #print(self.w)

        self.wm_w = self.watermark_resize.shape[1]
        self.wm_h = self.watermark_resize.shape[0]

        if pos_index == 1:
            self.LeftTop = [self.w * self.bleedingWidth, self.h * self.bleedingWidth]
            self.wm_x, self.wm_y = int(self.LeftTop[0]), int(self.LeftTop[1])
        elif pos_index == 2:
            self.MidTop = [(self.w - self.wm_w) / 2, self.h * self.bleedingWidth]
            self.wm_x, self.wm_y = int(self.MidTop[0]), int(self.MidTop[1])
        elif pos_index == 3:
            self.RightTop = [self.w * (1 - self.bleedingWidth) - self.wm_w, self.h * self.bleedingWidth]
            self.wm_x, self.wm_y = int(self.RightTop[0]), int(self.RightTop[1])
        elif pos_index == 4:
            self.LeftMid = [self.w * self.bleedingWidth, (self.h - self.wm_h) / 2]
            self.wm_x, self.wm_y = int(self.LeftMid[0]), int(self.LeftMid[1])
        elif pos_index == 5:
            self.Center = [(self.w - self.wm_w) / 2, (self.h - self.wm_h) / 2]
            self.wm_x, self.wm_y = int(self.Center[0]), int(self.Center[1])
        elif pos_index == 6:
            self.RightMid = [self.w * (1 - self.bleedingWidth) - self.wm_w, (self.h - self.wm_h) / 2]
            self.wm_x, self.wm_y = int(self.RightMid[0]), int(self.RightMid[1])
        elif pos_index == 7:
            self.LeftBottom = [self.w * self.bleedingWidth, self.h * (1 - self.bleedingWidth) - self.wm_h]
            self.wm_x, self.wm_y = int(self.LeftBottom[0]), int(self.LeftBottom[1])
        elif pos_index == 8:
            self.MidBottom = [(self.w - self.wm_w) / 2, self.h * (1 - self.bleedingWidth) - self.wm_h]
            self.wm_x, self.wm_y = int(self.MidBottom[0]), int(self.MidBottom[1])
        elif pos_index == 9:
            self.RightBottom = [self.w * (1 - self.bleedingWidth) - self.wm_w, self.h * (1 - self.bleedingWidth) - self.wm_h]
            self.wm_x, self.wm_y = int(self.RightBottom[0]), int(self.RightBottom[1])

        self.wm_left_x = self.wm_x
        self.wm_right_x = self.wm_x + self.wm_w
        self.wm_top_y = self.wm_y
        self.wm_bottom_y = self.wm_y + self.wm_h

        #接下来做切片操作，横坐标切片一次，纵坐标切片一次即可
        #水印如果是RGBA色彩空间的png图片的话，这里还不能直接在ROI上绘制图片

        #print(self.watermark_resize.shape)
        #print(self.image.shape)

        # 定义水印ROI: Region Of Interest
        # ROI左上角坐标(self.wm_x, self.wm_y)
        self.ROI = self.image[self.wm_top_y : self.wm_bottom_y, self.wm_left_x : self.wm_right_x]
        self.result_img = cv.addWeighted(self.ROI, 1, self.watermark_resize, opacity, 0)
        self.image[self.wm_top_y : self.wm_bottom_y, self.wm_left_x : self.wm_right_x] = self.result_img

        return self.image

if __name__ == '__main__':
    src = './Images/watermark_demo.jpg'
    watermark_image = './img/照片水印-白色-test.jpg'
    dst = WaterMark().watermark(src, watermark_image, 8, 0.75, 0.5)
    print(watermark_image)
    cv.imshow('Watermark Test', dst)
    cv.waitKey(0)
    cv.destroyAllWindows()